import os
import numpy as np
from PIL import Image
import joblib
import gradio as gr

LABEL_MAP = {0: 'cat', 1: 'dog'}
MODEL_PATH = os.path.join('models', 'best_model.joblib')


def preprocess_image_flat(pil_img: Image.Image) -> np.ndarray:
    # 转灰度并缩放至 32x32，然后展平并标准化到 [0,1]
    img = pil_img.convert('L').resize((32, 32))
    arr = np.asarray(img, dtype=np.float32) / 255.0
    return arr.reshape(1, -1)


def load_model():
    if not os.path.exists(MODEL_PATH):
        raise FileNotFoundError(f"未找到已训练模型: {MODEL_PATH}。请先运行 train_ensemble.py 保存最佳模型。")
    payload = joblib.load(MODEL_PATH)
    return payload['model']


MODEL = load_model()


def predict(img: Image.Image):
    x = preprocess_image_flat(img)
    # 优先使用 predict_proba，否则退化为 hard 预测
    proba = None
    if hasattr(MODEL, 'predict_proba'):
        try:
            proba = MODEL.predict_proba(x)[0]
        except Exception:
            proba = None
    if proba is None and hasattr(MODEL, 'decision_function'):
        # 将 decision_function 通过 softmax 近似为概率
        z = np.asarray(MODEL.decision_function(x), dtype=np.float32).reshape(-1)
        z = z - z.max()
        ez = np.exp(z)
        proba = ez / ez.sum()
    if proba is None:
        y = int(MODEL.predict(x)[0])
        proba = np.zeros(2, dtype=np.float32)
        proba[y] = 1.0

    # 输出 gr.Label 需要的字典 {label: score}
    return {LABEL_MAP[i]: float(p) for i, p in enumerate(proba)}


def app():
    title = "Dog vs Cat (flat + CPU)"
    description = "上传一张猫狗图片，模型将在 CPU 上（flat 特征）进行推理。"
    example_note = "请确保已运行 train_ensemble.py 训练并保存 models/best_model.joblib。"

    iface = gr.Interface(
        fn=predict,
        inputs=gr.Image(type='pil', label='上传图片'),
        outputs=gr.Label(num_top_classes=2, label='预测概率'),
        title=title,
        description=f"{description}\n{example_note}",
        allow_flagging='never'
    )
    iface.launch()


if __name__ == '__main__':
    app()
